import json
import sys

def generate_m3u(data):
    uhd_output = []
    uhd_output.append("#EXTM3U")
    
    other_output = []
    other_output.append("#EXTM3U")

    uhd_urls = set()
    other_urls = set()

    for channel in data["channels"]:
        channel_name = channel["title"]
        
        best_url = ""
        best_quality = -1 # 4K > HD > SD
        is_uhd = False
        
        for phychannel in channel.get("phychannels", []):
            quality = -1
            current_is_uhd = False
            if phychannel.get("bitrateTypeName") == "4K超高清":
                quality = 2
                current_is_uhd = True
            elif phychannel.get("bitrateTypeName") == "高清":
                quality = 1
            elif phychannel.get("bitrateTypeName") == "标清":
                quality = 0

            if quality > best_quality:
                if "hwurl" in phychannel.get("params", {}):
                    best_url = phychannel["params"]["hwurl"]
                    best_quality = quality
                    is_uhd = current_is_uhd
                elif "zteurl" in phychannel.get("params", {}):
                    best_url = phychannel["params"]["zteurl"]
                    best_quality = quality
                    is_uhd = current_is_uhd

        if not best_url:
            if "hwurl" in channel.get("params", {}):
                best_url = channel["params"]["hwurl"]
            elif "zteurl" in channel.get("params", {}):
                best_url = channel["params"]["zteurl"]

        if best_url:
            if is_uhd:
                if best_url not in uhd_urls:
                    uhd_output.append(f"#EXTINF:-1,{channel_name} (UHD)")
                    uhd_output.append(best_url)
                    uhd_urls.add(best_url)
            else:
                if best_url not in other_urls:
                    other_output.append(f"#EXTINF:-1,{channel_name}")
                    other_output.append(best_url)
                    other_urls.add(best_url)
    
    return "\n".join(other_output), "\n".join(uhd_output)

if __name__ == "__main__":
    with open("/Users/tiger/workspace/CMCC-IPTV-main/频道总信息/getAllChannel2.json", "r") as f:
        data = json.load(f)
    
    other_m3u_content, uhd_m3u_content = generate_m3u(data)

    if len(sys.argv) > 2:
        with open(sys.argv[1], "w") as f:
            f.write(other_m3u_content)
        with open(sys.argv[2], "w") as f:
            f.write(uhd_m3u_content)
    elif len(sys.argv) > 1:
        with open(sys.argv[1], "w") as f:
            f.write(other_m3u_content)
        # Default name for UHD playlist if only one argument is provided
        with open("uhd_playlist.m3u", "w") as f:
            f.write(uhd_m3u_content)
    else:
        print("--- Standard Playlist ---")
        print(other_m3u_content)
        print("\n--- UHD Playlist ---")
        print(uhd_m3u_content)
